import torch
from torch import nn
from .attention import Attention, MyAttentionLayer

# 只有最后一层有注意力机制
class bigru_self_attention_model(nn.Module):
    def __init__(self, configs):
        super(bigru_self_attention_model, self).__init__()
        skips = []
        self.bigru_block1 = nn.Sequential(
            nn.GRU(configs.input_channels, 9, 3, batch_first=True, bidirectional=True, dropout=0.1),
        )

        self.dropout1 = nn.Dropout(configs.dropout)
        self.attention_layer = Attention(configs.features_len,heads=1,dropout=0.1)
        # self.attention_layer = MyAttentionLayer(configs.final_out_channels)


        # self.bigru_block2 = nn.Sequential(
        #     # nn.BatchNorm1d(configs.input_channels),
        #     # torch.nn.GRU(input_dim, hidden_size, n_layers, batch_first=True, bidirectional=True)
        #     nn.GRU(configs.input_channels, 32, 1, batch_first=True, bidirectional=True),
        #     nn.Dropout(configs.dropout),
        #
        # )
        #
        # self.bigru_block3 = nn.Sequential(
        #     # nn.BatchNorm1d(configs.input_channels),
        #     # torch.nn.GRU(input_dim, hidden_size, n_layers, batch_first=True, bidirectional=True)
        #     nn.GRU(configs.input_channels, 32, 1, batch_first=True, bidirectional=True),
        #     # nn.Dropout(configs.dropout)
        # )

        model_output_dim = configs.features_len
        self.logits = nn.Linear(model_output_dim * configs.final_out_channels, configs.num_classes)

    def forward(self, x_in):
        output, hidden = self.bigru_block1(x_in)
        # x = self.liner_block1(x)
        # x = self.bigru_block2(x)
        # -------------------注意力层-----------------------------
        # print("output[0].shape:", output[0].shape)
        # print("hidden.shape:", hidden.shape)
        output  = self.attention_layer(output)
        # print("output[0].shape:", output[0].shape)
        # print("output.shape:", output.shape)
        # ------------------------------------------------
        x_flat = output.reshape(output.shape[0], -1)
        # print("x_flat.shape:", x_flat.shape)

        logits = self.logits(x_flat)
        return logits, output

    # def forward(self, x, h):
    #     out, h = self.bigru_block1(x, h)
    #     out = self.fc(self.relu(out[:,-1]))
    #     return out, h